package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os/exec"
	"strings"
	"time"

	"golang.org/x/text/encoding/simplifiedchinese"
)

var process string
var fileName string
var path string
var port string
var interval int
var count int
var timeout int

func init() {
	flag.StringVar(&process, "process", "", "进程名称")
	flag.StringVar(&fileName, "fileName", "", "进程实际文件名称(不填则默认使用进程名称)")
	flag.StringVar(&path, "path", "", "进程所在目录的绝对路径")
	flag.StringVar(&port, "port", "", "端口号(选填)")
	flag.IntVar(&interval, "interval", 60, "检测间隔(秒)")
	flag.IntVar(&count, "count", 5, "连续未响应次数")
	flag.IntVar(&timeout, "timeout", 10, "端口尝试连接时间(秒)")
	flag.Parse()

	if fileName == "" {
		fileName = process
	}
}

func main() {
	for {
		if processExist(process) {
			if processNameScan(process, count) || !processPortScan(port, timeout) {
				if processStop(process) {
					processStart(fileName, path)
				}
			}
		} else {
			processStart(fileName, path)
		}

		time.Sleep(time.Second * time.Duration(interval))
	}
}

// 扫描进程是否存在
func processExist(processName string) bool {
	cmd := fmt.Sprintf(`tasklist /FI "IMAGENAME eq %v"`, processName)
	log.Printf("进程扫描：开始执行 [%v]", cmd)
	out, err := exec.Command("powershell", cmd).Output()
	if err != nil {
		log.Println("进程扫描：执行失败", err)
		return true
	}
	decodeByte, _ := simplifiedchinese.GB18030.NewDecoder().Bytes(out)
	log.Println("进程扫描：执行结果", string(decodeByte))
	return strings.Contains(string(decodeByte), processName)
}

// 扫描进程是否未响应
func processNameScan(processName string, count int) bool {
	cmd := fmt.Sprintf(`tasklist /FI "STATUS eq NOT RESPONDING" /FI "IMAGENAME eq %v"`, processName)

	result := false
	for i := 1; i <= count; i++ {
		log.Printf("状态扫描：%v第%v次执行 [%v]", processName, i, cmd)

		out, err := exec.Command("powershell", cmd).Output()
		if err != nil {
			log.Println("状态扫描：执行失败", err)
			time.Sleep(time.Second * 5)
			continue
		}

		decodeByte, _ := simplifiedchinese.GB18030.NewDecoder().Bytes(out)
		result = strings.Contains(string(decodeByte), processName)
		if result {
			log.Println("状态扫描：状态[未响应]")
		} else {
			log.Println("状态扫描：状态[正常运行]")
			break
		}

		time.Sleep(time.Second * 5)
	}

	return result
}

// 扫描端口是否存在
func processPortScan(port string, timeout int) bool {
	if port == "" {
		return true
	}
	log.Println("端口扫描：开始执行")
	conn, err := net.DialTimeout("tcp", net.JoinHostPort("127.0.0.1", port), time.Second*time.Duration(timeout))
	if conn != nil {
		conn.Close()
	}
	log.Println("端口扫描：执行结果", err == nil)
	return err == nil
}

// 进程启动
func processStart(fileName string, path string) bool {
	cmd := fmt.Sprintf(`Start-Process -workingdirectory "%v" "%v"`, path, fileName)
	log.Printf("进程启动：开始执行 [%v]", cmd)
	out, err := exec.Command("powershell", cmd).Output()
	if err != nil {
		log.Println("进程启动：执行失败", err)
		return false
	}
	decodeByte, _ := simplifiedchinese.GB18030.NewDecoder().Bytes(out)
	log.Println("进程启动：执行结果", string(decodeByte))
	return true
}

// 进程关闭
func processStop(processName string) bool {
	cmd := fmt.Sprintf(`taskkill /F /T /IM "%v"`, processName)
	log.Printf("进程关闭：开始执行 [%v]", cmd)
	out, err := exec.Command("powershell", cmd).Output()
	if err != nil {
		log.Println("进程关闭：执行失败", err)
		return false
	}
	decodeByte, _ := simplifiedchinese.GB18030.NewDecoder().Bytes(out)
	log.Println("进程关闭：执行结果", string(decodeByte))
	return true
}
